package com.chen.handler;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.util.ObjUtil;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.chen.bean.constants.WxConstants;
import com.chen.bean.entity.DiaLogEntity;
import com.chen.bean.entity.MessageEntity;
import com.chen.bean.entity.UserEntity;
import com.chen.bean.vo.MessageRequest;
import com.chen.exception.BusinessException;
import com.chen.manager.QwenManager;
import com.chen.service.DialogService;
import com.chen.service.MessageService;
import com.chen.util.CacheUtils;
import com.chen.util.JsonMapper;
import com.fasterxml.jackson.core.type.TypeReference;
import io.reactivex.Flowable;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.checkerframework.checker.units.qual.A;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.util.ArrayList;
import java.util.Date;
import java.util.List;
import java.util.Objects;
import java.util.concurrent.atomic.AtomicInteger;

@Component
@Slf4j
public class WxChatWebSocketHandler extends TextWebSocketHandler {

    @Autowired
    private CacheUtils cacheUtils;

    @Autowired
    private MessageService messageService;

    @Autowired
    private DialogService dialogService;

    @Autowired
    private QwenManager qwenManager;


    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        MessageRequest messageRequest = JsonMapper.fromJson(message.getPayload(), MessageRequest.class);
        if (Objects.isNull(messageRequest)) return;
        String messageText = messageRequest.getMessage();
        String dialogNumber = messageRequest.getDialogNumber();
        session.getAttributes().put(WxConstants.WEB_SOCKET_DIALOG_NUMBER,dialogNumber);
       if (StringUtils.isBlank(messageText) || StringUtils.isBlank(dialogNumber)){
           throw new BusinessException("发送消息和对话编号不能为空");
       }
        List<MessageEntity> historyEntityList;
        //请求通义千问消息体
        Message userMsg = Message.builder().role(Role.USER.getValue()).content(messageText).build();
        //历史消息缓存key
        String messageSuffixCacheKey = WxConstants.WX_HISTORY_MESSAGE_CACHE + dialogNumber;

        String json = JsonMapper.toJson(cacheUtils.getCacheList(messageSuffixCacheKey));
        historyEntityList = JsonMapper.formJson(json, new TypeReference<List<MessageEntity>>() {
        });
        //缓存为空查询数据库
        if (CollectionUtil.isEmpty(historyEntityList)){
            historyEntityList =  messageService.queryByDialogNumber(dialogNumber);
        }
        List<Message> historyMessageList = messageService.buildMessageList(historyEntityList);
        List<Message> messageList = new ArrayList<>(historyMessageList);

        //本次用户本次所发消息
        MessageEntity messageEntity = messageService.buildMessageEntity(dialogNumber, userMsg.getRole(), userMsg.getContent());
        historyEntityList.add(messageEntity);
        messageList.add(userMsg);
        //调用通义千问
        Flowable<GenerationResult> flowable = qwenManager.qwenStreamCall(messageList);
        StringBuffer sb = new StringBuffer();
        //本次消息输出总长度
        AtomicInteger outputTokens = new AtomicInteger(0);
        MessageEntity assistantEntity = new MessageEntity();
        List<MessageEntity> finalHistoryEntityList = historyEntityList;
        flowable.subscribe(
                result -> {
                    // 处理每个 GenerationResult
                    Message assistantMsg = result.getOutput().getChoices().get(0).getMessage();
                    String content = assistantMsg.getContent();
                    outputTokens.addAndGet(result.getUsage().getOutputTokens());
                    assistantEntity.setRole(assistantMsg.getRole());
                    assistantEntity.setInputTokens(result.getUsage().getInputTokens());
                    sb.append(content);
                    session.sendMessage(new TextMessage(content));
                },
                throwable -> {
                    // 处理错误
                    log.error("Error in stream call: {}", throwable.getMessage());
                },
                () -> {
                    // 完成时的回调
                    log.info("本轮对话完成。。");
                    assistantEntity.setOutputTokens(outputTokens.get());
                    Integer totalTokens =assistantEntity.getInputTokens() + assistantEntity.getOutputTokens();
                    assistantEntity.setTotalTokens(totalTokens);
                    assistantEntity.setDialogNumber(dialogNumber);
                    assistantEntity.setContent(sb.toString());
                    assistantEntity.setCreateTime(new Date());
                    assistantEntity.setUpdateTime(new Date());
                    finalHistoryEntityList.add(assistantEntity);
                    //消息内容暂存redis,异步执行
                    messageService.saveMessageToRedis(messageSuffixCacheKey,finalHistoryEntityList);
                }
        );

    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        UserEntity user = (UserEntity)session.getAttributes().get(WxConstants.WX_WEBSOCKET_ATTRIBUTE);
        if (Objects.isNull(user)){
            log.error("用户信息为空，请排查日志");
            return;
        }
        log.info("用户{}建立连接",user.getNickname());
        String message = "久违了，"+user.getNickname()+"，我已等你许久";
        session.sendMessage(new TextMessage(message));
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        log.info("websocket连接出错");
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        UserEntity user = (UserEntity)session.getAttributes().get(WxConstants.WX_WEBSOCKET_ATTRIBUTE);
        if (Objects.isNull(user)){
            log.error("用户信息为空，请排查日志");
            return;
        }
        String dialogNumber = (String) session.getAttributes().get(WxConstants.WEB_SOCKET_DIALOG_NUMBER); // 获取对话编号
        if (StringUtils.isBlank(dialogNumber)){
            log.error("对话编号为空，请排查日志");
            return;
        }
        String key = WxConstants.WX_HISTORY_MESSAGE_CACHE + dialogNumber;
        String json = JsonMapper.toJson(cacheUtils.getCacheList(key));
        List<MessageEntity> cacheList = JsonMapper.formJson(json, new TypeReference<>() {
        });
        //如果缓存为空，说明未发生聊天，直接return
        if (CollectionUtil.isEmpty(cacheList)) return;
        //保存消息列表message
        //存消息记录 去重已经存入的，id不为空说明已入库
        List<MessageEntity> distinct = cacheList.stream().filter(entity -> ObjUtil.isNull(entity.getId())).toList();
        if (CollectionUtil.isNotEmpty(distinct)) {
            log.info("用户：{} 结束本轮对话，新增消息数量：{}",user.getNickname(),distinct.size());
            messageService.batchSave(distinct);
        }
        DiaLogEntity diaLogEntity = dialogService.queryEntityByDialogNumber(dialogNumber);
        if (ObjUtil.isNull(diaLogEntity)){
            //为空说明新对话
            String content = cacheList.get(0).getContent();
            // 取前十位字符防止超出数据库长度
            String dialogName = content.length() > 10 ? content.substring(0, 10) : content;
            dialogService.save(dialogNumber, dialogName,user.getId());
        }else {
            //不为空更新
            diaLogEntity.setUpdateTime(new Date());
            dialogService.update(diaLogEntity);
        }
        //清除缓存
        cacheUtils.deleteObject(key);
    }
}
